from numpy import extract
import pandas as pd
from collections import defaultdict
import glob
import json
import os
from tqdm import tqdm

TEXT_LEN = 11

audioset_csvs = sorted(glob.glob("../src/AudioSet/meta/*.csv"))
mapping_file = "../src/DESED_AS_label_map.json"
desed_filenames = "../src/DESED_all_filenames.json"

with open(mapping_file, "r") as f:
    label_map = json.load(f)    

# reverse the label map
reverse_label_map = defaultdict(str)
for key, value in label_map.items():
    for v in value:
        reverse_label_map[v] = key

label_map = reverse_label_map
valid_labels = set(label_map.values())
# load all desed filenames
with open(desed_filenames, "r") as f:
    desed_filenames = set(json.load(f))

# load audioset csvs
extracted_df = pd.DataFrame(columns=["# YTID", "start_seconds", "end_seconds", "positive_labels", "desed_labels"])
for csv in audioset_csvs:
    print(f"Processing {csv}...")
    df = pd.read_csv(csv, skiprows=2, sep=", ", engine='python')
    df["positive_labels"] = df["positive_labels"].apply(lambda x: x[1:-1].replace(" ", ""))
    df["desed_filename"] = df.apply(lambda row: f"Y{row['# YTID']}_{row['start_seconds']:.3f}_{row['end_seconds']:.3f}", axis=1)
    df["desed_labels"] = df["positive_labels"].apply(lambda x: ",".join([label_map.get(label) if label_map.get(label) is not None else "other" for label in x.split(",")]))
    # print valid rows
    
    df_valid = df[df["desed_labels"].apply(lambda x: any([label in valid_labels for label in x.split(",")]))]
    print(f"Total rows: {len(df)}, Valid rows: {len(df_valid)}")
    extracted_df = pd.concat([extracted_df, df_valid[["# YTID", "start_seconds", "end_seconds", "positive_labels", "desed_labels", "desed_filename"]]], ignore_index=True)
    # print(df_valid.head())
# save the extracted dataframe to csv
os.makedirs("../src/extracted/", exist_ok=True)
extracted_df.to_csv("../src/extracted/DESED_Audioset_label_mapping.csv", index=False)

# find if desed file appeared in the csv
# valid_labels = set()
# for csv in audioset_csvs:
#     df = pd.read_csv(csv, skiprows=2, sep=", ", engine='python')
#     print(csv, len(df))
#     df["desed_filename"] = df.apply(lambda row: f"Y{row['# YTID']}_{row['start_seconds']:.3f}_{row['end_seconds']:.3f}", axis=1)
#     as_filenames = set(df["desed_filename"].tolist())
#     intersec = as_filenames.intersection(desed_filenames)
#     if intersec:
#         print(f"Found {len(intersec)} desed filenames in {csv}.")
#         for filename in intersec:
#             labels = df[df["desed_filename"] == filename]["positive_labels"].item()[1:-1]
#             valid_labels.update(set(labels.split(",")))
#         print("Contained labels:", len(valid_labels))
#     else:
#         print(f"No desed filenames found in {csv}.")
    